import tensorflow as tf
import numpy as np
# import IPython.display as display
import PIL.Image
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image

# 图像标准化
def normalize_image(img):
    img = 255 * (img + 1.0) / 2.0
    return tf.cast(img, tf.uint8)

# 图像可视化
def show_image(img):
    plt.imshow(img)
    plt.show()
    # display.display(PIL.Image.fromarray(np.array(img)))

# 保存图像文件
def save_image(img, file_name):
    PIL.Image.fromarray(np.array(img)).save(file_name)


# ======== 数据预处理 ===========
# 定义图像噪声
img_noise = np.random.uniform(size=(300, 300, 3)) + 100.0
img_noise = img_noise.astype(np.float32)
show_image(normalize_image(img_noise))

img = tf.keras.applications.inception_v3.preprocess_input(img_noise)
img = tf.convert_to_tensor(img)

# ======== 预训练模型的加载 ===========
# 导入ImageNet数据集的图像识别预训练InceptionV3模型
# 去掉顶层，这样能接受新的训练数据shape
base_model = tf.keras.applications.InceptionV3(include_top=False, weights='imagenet')
base_model.summary()

# ======== 选择卷积层和通道 ===========
# 最大限度地激活这些层的指定的层
layer_names = ['mixed3', 'mixed5']
layers = [base_model.get_layer(name).output for name in layer_names]
print(layers)

# 创建特征提取模型
dream_model = tf.keras.Model(inputs=base_model.input, outputs=layers)
dream_model.summary()

# ======== 计算损失 ===========
# 损失是选中层的通道输出
def calc_loss(img, model):
    # 对图像做变形，由（300,300,3）扩展为（1,300,300,3）
    img = tf.expand_dims(img, axis=0)
    # 图像通过模型前向传播得到计算结果
    layer_activations = model(img)

    losses = []
    # 计算选中每层的计算结果均值
    for act in layer_activations:
        loss = tf.math.reduce_mean(act)
        losses.append(loss)
    # 返回选中每通道的计算结果和
    return tf.reduce_sum(losses)


# 定义图像优化过程
def render_deepdream(model, img, steps=100, step_size=0.01):
    for n in range(steps):
        with tf.GradientTape() as tape:
            # 对img进行梯度变换
            tape.watch(img)
            loss = calc_loss(img, model)
        # 计算损失相对于输入图像像素的梯度
        gradients = tape.gradient(loss, img)
        # 归一化梯度值
        gradients /= tf.math.reduce_std(gradients) + 1e-8
        # 在梯度上升中，损失值越来越大，因此可以直接添加损失值到图像中，因为它们的shape相同
        img = img + gradients * step_size
        img = tf.clip_by_value(img, -1, 1)
        # 输出过程提示信息
        if (n + 1) % 10 == 0:
            print('Step {}/{}, loss {}'.format(n+1, steps, loss))
    return img

# ======== 开始做梦 ===========
# 调用优化过程
dream_img = render_deepdream(dream_model, img)
# 标准化图像
dream_img = normalize_image(dream_img)
# 显示结果图像
show_image(dream_img)
# 保存结果图像
file_name = 'deepdream_{}.jpg'.format(layer_names)
save_image(dream_img, file_name)